import akshare as ak
import pandas as pd
import numpy as np
from datetime import datetime, date, timedelta
from sqlalchemy import create_engine, text
import time
from functools import wraps
from db_config import DBConfig

# 简单的重试装饰器实现
def retry(max_attempts=5, delay=1, backoff=2, exceptions=(Exception,)):
    """
    简单的重试装饰器
    :param max_attempts: 最大重试次数
    :param delay: 初始延迟时间（秒）
    :param backoff: 退避系数
    :param exceptions: 需要重试的异常类型
    """
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            attempts = 0
            current_delay = delay
            while attempts < max_attempts:
                try:
                    return func(*args, **kwargs)
                except exceptions as e:
                    attempts += 1
                    if attempts == max_attempts:
                        print(f"函数 {func.__name__} 重试 {max_attempts} 次后仍然失败: {e}")
                        raise
                    print(f"函数 {func.__name__} 第 {attempts} 次失败: {e}, {current_delay}秒后重试...")
                    time.sleep(current_delay)
                    current_delay *= backoff  # 指数退避
        return wrapper
    return decorator

def clean_stock_code(stock_code):
    """
    清理股票代码，移除.sz和.sh后缀
    :param stock_code: 原始股票代码
    :return: 清理后的股票代码
    """
    return stock_code.replace('.SZ', '').replace('.SH', '').replace('.BJ', '')

def sync_margin_summary_to_db(start_date=None, end_date=None):
    """
    同步融资融券汇总数据到数据库
    """
    try:
        # 获取融资融券汇总数据
        raw_df = get_margin_summary_data(start_date, end_date)
        if raw_df.empty:
            return {
                "success": True,
                "message": "未获取到融资融券汇总数据",
                "records": 0
            }

        # 预处理汇总数据
        df = prepare_margin_summary_data(raw_df)
        if df.empty:
            return {
                "success": True,
                "message": "融资融券汇总数据处理后无有效数据",
                "records": 0
            }

        # 插入数据库
        records = insert_margin_summary_to_mysql(df)

        return {
            "success": True,
            "message": f"成功同步融资融券汇总数据 {records} 条",
            "records": records
        }
    except Exception as e:
        error_msg = f"融资融券汇总数据处理失败: {e}"
        return {
            "success": False,
            "message": error_msg,
            "records": 0
        }

def sync_margin_detail_to_db(trade_date):
    """
    同步指定日期的融资融券明细数据到数据库
    """
    try:
        # 获取融资融券明细数据
        raw_df = get_margin_detail_data(trade_date)
        if raw_df.empty:
            return {
                "success": True,
                "message": f"{trade_date} 未获取到融资融券明细数据",
                "records": 0
            }

        # 预处理明细数据
        df = prepare_margin_detail_data(raw_df)
        if df.empty:
            return {
                "success": True,
                "message": f"{trade_date} 融资融券明细数据处理后无有效数据",
                "records": 0
            }

        # 插入数据库
        records = insert_margin_detail_to_mysql(df)

        return {
            "success": True,
            "message": f"成功同步 {trade_date} 融资融券明细数据 {records} 条",
            "records": records
        }
    except Exception as e:
        error_msg = f"融资融券明细数据处理失败: {e}"
        return {
            "success": False,
            "message": error_msg,
            "records": 0
        }

@retry(max_attempts=5, delay=1, backoff=2, exceptions=(Exception,))
def get_margin_summary_data(start_date=None, end_date=None):
    """
    获取融资融券汇总数据（带重试机制）
    """
    try:
        df = ak.stock_margin_sse()
        return df
    except Exception as e:
        print(f"获取融资融券汇总数据异常: {e}")
        return pd.DataFrame()

@retry(max_attempts=5, delay=1, backoff=2, exceptions=(Exception,))
def get_margin_detail_data(trade_date):
    """
    获取指定日期的融资融券明细数据（带重试机制）
    """
    try:
        # 将日期格式转换为akshare需要的格式
        if isinstance(trade_date, date):
            date_str = trade_date.strftime("%Y%m%d")
        elif isinstance(trade_date, str):
            # 如果是YYYY-MM-DD格式，转换为YYYYMMDD
            if len(trade_date) == 10:
                date_str = trade_date.replace("-", "")
            else:
                date_str = trade_date
        else:
            date_str = trade_date

        df = ak.stock_margin_detail_sse(date=date_str)
        return df
    except Exception as e:
        print(f"获取 {trade_date} 融资融券明细数据异常: {e}")
        return pd.DataFrame()

def prepare_margin_summary_data(raw_df):
    """
    按 stock_margin_summary 表字段要求预处理汇总数据
    """
    if raw_df.empty:
        return pd.DataFrame()

    # 复制数据避免修改原数据
    df = raw_df.copy()

    # 基于列的索引位置重命名列（更可靠的方法）
    # 融资融券汇总数据列顺序：交易日期, 融资余额, 融资买入额, 融券余额, 融券卖出额, 融券余额, 融资融券余额
    if len(df.columns) >= 7:
        df.columns = ['trade_date', 'financing_balance', 'financing_buy_amount',
                     'securities_lending_balance', 'securities_lending_amount',
                     'securities_lending_balance_dup', 'total_margin_balance']

        # 删除重复的融券余额列
        if 'securities_lending_balance_dup' in df.columns:
            df = df.drop('securities_lending_balance_dup', axis=1)

    # 只保留需要的列
    expected_columns = ['trade_date', 'financing_balance', 'securities_lending_balance',
                       'total_margin_balance', 'financing_buy_amount', 'securities_lending_amount']
    available_columns = [col for col in expected_columns if col in df.columns]
    df = df[available_columns]

    # 处理日期字段
    if 'trade_date' in df.columns:
        df['trade_date'] = pd.to_datetime(df['trade_date'], errors='coerce').dt.date

    # 处理数值字段，将无穷大和NaN替换为None
    numeric_columns = df.select_dtypes(include=[np.number]).columns.tolist()
    for col in numeric_columns:
        df[col] = df[col].replace([np.inf, -np.inf], np.nan)
        df[col] = df[col].where(pd.notnull(df[col]), None)

        # 对于金额字段，保留2位小数
        if col in ['financing_balance', 'securities_lending_balance', 'total_margin_balance',
                  'financing_buy_amount', 'securities_lending_amount', 'total_margin_amount']:
            df[col] = df[col].apply(lambda x: round(x, 2) if x is not None else None)

    return df

def prepare_margin_detail_data(raw_df):
    """
    按 stock_margin_detail 表字段要求预处理明细数据
    """
    if raw_df.empty:
        return pd.DataFrame()

    # 复制数据避免修改原数据
    df = raw_df.copy()

    # 基于列的索引位置重命名列（更可靠的方法）
    # 实际数据结构：交易日期, 证券代码, 证券简称, 融资余额, 融资买入额, 融券余额, 融券卖出额, 融资余额(重复), 融券余额(重复)
    if len(df.columns) >= 7:
        # 只取前7列，有效数据列
        df = df.iloc[:, :7]
        df.columns = ['trade_date', 'stock_code', 'stock_name', 'financing_balance',
                     'financing_buy_amount', 'securities_lending_balance', 'securities_lending_amount']

        # 计算融资融券余额（融资余额 + 融券余额）
        df['total_margin_balance'] = df['financing_balance'] + df['securities_lending_balance']
        df['total_margin_amount'] = df['financing_buy_amount'] + df['securities_lending_amount']

    # 只保留需要的列
    expected_columns = ['trade_date', 'stock_code', 'stock_name', 'financing_balance',
                       'securities_lending_balance', 'total_margin_balance', 'financing_buy_amount',
                       'securities_lending_amount', 'total_margin_amount']
    available_columns = [col for col in expected_columns if col in df.columns]
    df = df[available_columns]

    # 处理日期字段
    if 'trade_date' in df.columns:
        df['trade_date'] = pd.to_datetime(df['trade_date'], errors='coerce').dt.date

    # 清理股票代码
    if 'stock_code' in df.columns:
        df['stock_code'] = df['stock_code'].astype(str).apply(clean_stock_code)

    # 处理数值字段，将无穷大和NaN替换为None
    numeric_columns = df.select_dtypes(include=[np.number]).columns.tolist()
    for col in numeric_columns:
        df[col] = df[col].replace([np.inf, -np.inf], np.nan)
        df[col] = df[col].where(pd.notnull(df[col]), None)

        # 对于金额字段，保留2位小数
        if col in ['financing_balance', 'securities_lending_balance', 'total_margin_balance',
                  'financing_buy_amount', 'securities_lending_amount', 'total_margin_amount']:
            df[col] = df[col].apply(lambda x: round(x, 2) if x is not None else None)

    return df

@retry(max_attempts=3, delay=1, backoff=2, exceptions=(Exception,))
def insert_margin_summary_to_mysql(df):
    """
    插入融资融券汇总数据到 MySQL（支持幂等）
    """
    db_config = DBConfig.config
    engine = create_engine(
        f"mysql+pymysql://{db_config['user']}:{db_config['password']}@{db_config['host']}:{db_config['port']}/{db_config['database']}?charset={db_config['charset']}"
    )

    try:
        # 获取数据库连接
        with engine.begin() as conn:
            for idx, row in df.iterrows():
                # 构建字典参数
                row_dict = {k: (None if pd.isna(v) else v) for k, v in row.to_dict().items()}

                # 构建列和占位符
                columns = ', '.join(row_dict.keys())
                placeholders = ', '.join([f":{k}" for k in row_dict.keys()])

                # 构建更新部分（排除主键和唯一键）
                update_clause = ', '.join([f"{col}=VALUES({col})" for col in row_dict.keys()])

                sql = f"""
                    INSERT INTO stock_margin_summary ({columns})
                    VALUES ({placeholders})
                    ON DUPLICATE KEY UPDATE {update_clause}
                """

                conn.execute(text(sql), row_dict)

        return len(df)
    except Exception as e:
        print(f"插入融资融券汇总数据失败: {e}")
        raise
    finally:
        engine.dispose()

@retry(max_attempts=3, delay=1, backoff=2, exceptions=(Exception,))
def insert_margin_detail_to_mysql(df):
    """
    插入融资融券明细数据到 MySQL（支持幂等）
    """
    db_config = DBConfig.config
    engine = create_engine(
        f"mysql+pymysql://{db_config['user']}:{db_config['password']}@{db_config['host']}:{db_config['port']}/{db_config['database']}?charset={db_config['charset']}"
    )

    try:
        # 获取数据库连接
        with engine.begin() as conn:
            for idx, row in df.iterrows():
                # 构建字典参数
                row_dict = {k: (None if pd.isna(v) else v) for k, v in row.to_dict().items()}

                # 构建列和占位符
                columns = ', '.join(row_dict.keys())
                placeholders = ', '.join([f":{k}" for k in row_dict.keys()])

                # 构建更新部分（排除主键和唯一键）
                update_clause = ', '.join([f"{col}=VALUES({col})" for col in row_dict.keys()])

                sql = f"""
                    INSERT INTO stock_margin_detail ({columns})
                    VALUES ({placeholders})
                    ON DUPLICATE KEY UPDATE {update_clause}
                """

                conn.execute(text(sql), row_dict)

        return len(df)
    except Exception as e:
        print(f"插入融资融券明细数据失败: {e}")
        raise
    finally:
        engine.dispose()

# -------------------- 批量同步函数 --------------------
def batch_sync_margin_detail(date_list):
    """
    批量同步融资融券明细数据
    :param date_list: 日期列表，格式为 ['2024-08-01', '2024-08-02', ...]
    """
    total_records = 0
    success_count = 0
    fail_count = 0

    for trade_date in date_list:
        try:
            result = sync_margin_detail_to_db(trade_date)
            if result["success"]:
                success_count += 1
                total_records += result["records"]
                if result["records"] > 0:
                    print(f"成功同步 {trade_date} 融资融券明细数据 {result['records']} 条")
                else:
                    print(f"{trade_date} 无新数据可同步")
            else:
                fail_count += 1
                print(f"同步 {trade_date} 融资融券明细数据失败: {result['message']}")
        except Exception as e:
            fail_count += 1
            print(f"同步 {trade_date} 融资融券明细数据异常: {e}")

    # 打印统计信息
    print("\n同步完成统计:")
    print(f"成功同步日期数: {success_count}")
    print(f"失败日期数: {fail_count}")
    print(f"总记录数: {total_records}")

# 使用示例
if __name__ == "__main__":
    # 同步融资融券汇总数据
    print("开始同步融资融券汇总数据...")
    result = sync_margin_summary_to_db()
    print(result)

    # 同步指定日期的融资融券明细数据
    print("\n开始同步融资融券明细数据...")
    result = sync_margin_detail_to_db("20240801")
    print(result)

    print('\n融资融券数据同步完成')